// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/core/buffer.h>

#include <algorithm>
#include <cstdint>
#include <utility>

#include <nebula/core/memory_pool_internal.h>

#include <turbo/utility/status.h>
#include <nebula/bits/bit_util.h>
#include <turbo/log/logging.h>
#include <nebula/util/slice_util_internal.h>
#include <nebula/util/string.h>

namespace nebula {

    turbo::Result<std::shared_ptr<Buffer>> Buffer::copy_slice(const int64_t start,
                                                             const int64_t nbytes,
                                                             MemoryPool *pool) const {
        // Sanity checks
        KCHECK_LE(start, size_);
        KCHECK_LE(nbytes, size_ - start);
                DKCHECK_GE(nbytes, 0);

        TURBO_MOVE_OR_RAISE(auto new_buffer, allocate_resizable_buffer(nbytes, pool));
        std::memcpy(new_buffer->mutable_data(), data() + start, static_cast<size_t>(nbytes));
        // R build with openSUSE155 requires an explicit shared_ptr construction
        return std::shared_ptr<Buffer>(std::move(new_buffer));
    }

    Buffer::Buffer() : Buffer(memory_pool::internal::kZeroSizeArea, 0) {}

    namespace {

        turbo::Status CheckBufferSlice(BufferSpan buffer, int64_t offset, int64_t length) {
            return internal::CheckSliceParams(buffer.size(), offset, length, "buffer");
        }

        turbo::Status CheckBufferSlice(BufferSpan buffer, int64_t offset) {
            if (TURBO_UNLIKELY(offset < 0)) {
                // Avoid UBSAN in subtraction below
                return turbo::out_of_range_error("Negative buffer slice offset");
            }
            return CheckBufferSlice(buffer, offset, buffer.size() - offset);
        }

    }  // namespace

    turbo::Result<std::shared_ptr<Buffer>> SliceBufferSafe(const std::shared_ptr<Buffer> &buffer,
                                                           int64_t offset) {
        TURBO_RETURN_NOT_OK(CheckBufferSlice(*buffer, offset));
        return SliceBuffer(buffer, offset);
    }

    turbo::Result<std::shared_ptr<Buffer>> SliceBufferSafe(const std::shared_ptr<Buffer> &buffer,
                                                           int64_t offset, int64_t length) {
        TURBO_RETURN_NOT_OK(CheckBufferSlice(*buffer, offset, length));
        return SliceBuffer(buffer, offset, length);
    }

    turbo::Result<std::shared_ptr<Buffer>> SliceMutableBufferSafe(
            const std::shared_ptr<Buffer> &buffer, int64_t offset) {
        TURBO_RETURN_NOT_OK(CheckBufferSlice(*buffer, offset));
        return SliceMutableBuffer(buffer, offset);
    }

    turbo::Result<std::shared_ptr<Buffer>> SliceMutableBufferSafe(
            const std::shared_ptr<Buffer> &buffer, int64_t offset, int64_t length) {
        TURBO_RETURN_NOT_OK(CheckBufferSlice(*buffer, offset, length));
        return SliceMutableBuffer(buffer, offset, length);
    }

    std::string Buffer::ToHexString() {
        return HexEncode(data(), static_cast<size_t>(size()));
    }

    bool Buffer::equals(const Buffer &other, const int64_t nbytes) const {
        return this == &other || (size_ >= nbytes && other.size_ >= nbytes &&
                                  (data_ == other.data_ ||
                                   !memcmp(data_, other.data_, static_cast<size_t>(nbytes))));
    }

    bool Buffer::equals(const Buffer &other) const {
        return this == &other || (size_ == other.size_ &&
                                  (data_ == other.data_ ||
                                   !memcmp(data_, other.data_, static_cast<size_t>(size_))));
    }

    std::string Buffer::to_string() const {
        return std::string(reinterpret_cast<const char *>(data_), static_cast<size_t>(size_));
    }

    void Buffer::CheckMutable() const { DKCHECK(is_mutable()) << "buffer not mutable"; }

    void Buffer::CheckCPU() const {
                DKCHECK(is_cpu()) << "not a CPU buffer (device: " << device()->to_string() << ")";
    }

    turbo::Result<std::shared_ptr<io::RandomAccessFile>> Buffer::get_reader(
            std::shared_ptr<Buffer> buf) {
        return buf->memory_manager_->GetBufferReader(std::move(buf));
    }

    turbo::Result<std::shared_ptr<io::OutputStream>> Buffer::get_writer(std::shared_ptr<Buffer> buf) {
        if (!buf->is_mutable()) {
            return turbo::invalid_argument_error("Expected mutable buffer");
        }
        return buf->memory_manager_->GetBufferWriter(std::move(buf));
    }

    turbo::Result<std::shared_ptr<Buffer>> Buffer::copy(std::shared_ptr<Buffer> source,
                                                        const std::shared_ptr<MemoryManager> &to) {
        return MemoryManager::CopyBuffer(source, to);
    }

    turbo::Result<std::unique_ptr<Buffer>> Buffer::copy_non_owned(
            const Buffer &source, const std::shared_ptr<MemoryManager> &to) {
        return MemoryManager::copy_non_owned(source, to);
    }

    turbo::Result<std::shared_ptr<Buffer>> Buffer::View(std::shared_ptr<Buffer> source,
                                                        const std::shared_ptr<MemoryManager> &to) {
        return MemoryManager::ViewBuffer(source, to);
    }

    turbo::Result<std::shared_ptr<Buffer>> Buffer::ViewOrCopy(
            std::shared_ptr<Buffer> source, const std::shared_ptr<MemoryManager> &to) {
        auto maybe_buffer = MemoryManager::ViewBuffer(source, to);
        if (maybe_buffer.ok()) {
            return maybe_buffer;
        }
        return MemoryManager::CopyBuffer(source, to);
    }

    class StlStringBuffer : public Buffer {
    public:
        explicit StlStringBuffer(std::string data) : input_(std::move(data)) {
            if (!input_.empty()) {
                data_ = reinterpret_cast<const uint8_t *>(input_.c_str());
                size_ = static_cast<int64_t>(input_.size());
                capacity_ = size_;
            }
        }

    private:
        std::string input_;
    };

    std::shared_ptr<Buffer> Buffer::from_string(std::string data) {
        return std::make_shared<StlStringBuffer>(std::move(data));
    }

    std::shared_ptr<Buffer> SliceMutableBuffer(const std::shared_ptr<Buffer> &buffer,
                                               const int64_t offset, const int64_t length) {
        return std::make_shared<MutableBuffer>(buffer, offset, length);
    }

    MutableBuffer::MutableBuffer(const std::shared_ptr<Buffer> &parent, const int64_t offset,
                                 const int64_t size)
            : MutableBuffer(reinterpret_cast<uint8_t *>(parent->mutable_address()) + offset,
                            size) {
                DKCHECK(parent->is_mutable()) << "Must pass mutable buffer";
        parent_ = parent;
    }

    turbo::Result<std::shared_ptr<Buffer>> allocate_bitmap(int64_t length, MemoryPool *pool) {
        TURBO_MOVE_OR_RAISE(auto buf, allocate_buffer(bit_util::BytesForBits(length), pool));
        // Zero out any trailing bits
        if (buf->size() > 0) {
            buf->mutable_data()[buf->size() - 1] = 0;
        }
        // R build with openSUSE155 requires an explicit shared_ptr construction
        return std::shared_ptr<Buffer>(std::move(buf));
    }

    turbo::Result<std::shared_ptr<Buffer>> allocate_empty_bitmap(int64_t length, MemoryPool *pool) {
        return allocate_empty_bitmap(length, kDefaultBufferAlignment, pool);
    }

    turbo::Result<std::shared_ptr<Buffer>> allocate_empty_bitmap(int64_t length, int64_t alignment,
                                                               MemoryPool *pool) {
        TURBO_MOVE_OR_RAISE(auto buf,
                            allocate_buffer(bit_util::BytesForBits(length), alignment, pool));
        memset(buf->mutable_data(), 0, static_cast<size_t>(buf->size()));
        // R build with openSUSE155 requires an explicit shared_ptr construction
        return std::shared_ptr<Buffer>(std::move(buf));
    }

    turbo::Status allocate_empty_bitmap(int64_t length, std::shared_ptr<Buffer> *out) {
        return allocate_empty_bitmap(length).try_value(out);
    }

    turbo::Result<std::shared_ptr<Buffer>> concatenate_buffers(
            const std::vector<std::shared_ptr<Buffer>> &buffers, MemoryPool *pool) {
        int64_t out_length = 0;
        for (const auto &buffer: buffers) {
            out_length += buffer->size();
        }
        TURBO_MOVE_OR_RAISE(auto out, allocate_buffer(out_length, pool));
        auto out_data = out->mutable_data();
        for (const auto &buffer: buffers) {
            // Passing nullptr to std::memcpy is undefined behavior, so skip empty buffers
            if (buffer->size() != 0) {
                std::memcpy(out_data, buffer->data(), buffer->size());
                out_data += buffer->size();
            }
        }
        // R build with openSUSE155 requires an explicit shared_ptr construction
        return std::shared_ptr<Buffer>(std::move(out));
    }

}  // namespace nebula
